import logging
from langchain_community.llms import Ollama
from langchain_core.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser

class QueryGenerator:
    def __init__(self):
        self.llm = Ollama(model="qwen2.5:latest")
        self.logger = logging.getLogger(__name__)
        self.sql_prompt = PromptTemplate.from_template(
            """您是一位SQL专家，帮助用户将自然语言转换为SQL查询。
            数据库名称：{database_name}
            数据库结构信息：
            {db_schema}
            用户查询：{user_query}
            请生成一个回答用户查询的SQL语句。只返回SQL语句，不要提供任何解释或其他内容。
            确保SQL语句可以在MySQL数据库中运行，并适当使用聚合函数、分组和过滤条件。"""
        )
        self.parser = StrOutputParser()
    
    def generate_sql(self, user_query, db_schema, database_name):
        try:
            self.logger.info(f"Generating SQL for query: {user_query}")
            chain = self.sql_prompt | self.llm | self.parser
            sql_query = chain.invoke({
                "user_query": user_query,
                "db_schema": db_schema,
                "database_name": database_name
            })
            sql_query = sql_query.replace("```sql", "").replace("```", "").strip()
            self.logger.info(f"Generated SQL: {sql_query}")
            return sql_query
        except Exception as e:
            self.logger.error(f"Error generating SQL: {str(e)}")
            raise